
/* Current Code Todos

- There is a race condition where the main thread calls P() on the semaphore
  and the sender thread immediately tries to send the data before the main thread
  copied the data into place. This can cause a segfault.

- Check how the window size is being updated.

- Handle the case of abrupt closes?
- Send SIGPIPE errors

- Fix valgrind uninitialized value errors for printf statements?


 */





#include "lab6.h"

packet* inFlightPackets[MAXWINDOW];
mutex closeLock;
mutex bufferAdd;
semaphore bufferEmpty;
semaphore bufferFull;
pthread_t senderThread;

int maxWindowSize;
int closed;
int growthRate;

//#define DEBUG 1

#define SLOWSTART (10)
#define AIMD ( 11 )


int my_socket( int domain, int type, int protocol ) {
  nextSlot = 0;
  startSlot = 0;
  currentTimeout = 1400;
  currentRTT = 1000;
  currentVar = 100;
  windowSize = 1;
  recipientClosed = 0;
  recvBase = 0;
  receivedFin = 0;
  sequenceNum = 0;
  maxWindowSize = 0;
  closed = 0;
  state = NORMAL;
  growthRate = SLOWSTART;

  memset(inFlightPackets, 0, MAXWINDOW * sizeof(packet*) );
  
  pthread_mutex_init( &closeLock, NULL );
  pthread_mutex_init( &bufferAdd, NULL );
  sem_init( &bufferEmpty, 0, MAXWINDOW );
  sem_init( &bufferFull, 0, 0 );
  uint64_t sock = socket(domain, type, protocol);
  
  

  return sock;

}

void my_send(int sock, void *buf, size_t len) {

  

  packet * p = Malloc (sizeof(packet));

  memset ( p, 0, sizeof(packet) );

  p->sequenceNumber = sequenceNum;

  //memset( p->data, 0, MAX_PACKET * sizeof(char) );

  sequenceNum += len;
  
  struct lab6_header * hdr = (struct lab6_header *) p->data;
  hdr->packet_num = htons( nextSlot );
  hdr->sequence_num = htonl(p->sequenceNumber);
  hdr->ack_flag = htonl ( 0 );
  hdr->fin_flag = htons ( 0 );
  hdr->length = htons( len );
  
  if ( buf ) {
    memcpy( hdr+1, buf, len );
  }
  else {
    hdr->fin_flag = htons ( 1 );
  }

  p->length = len;
  p->ack = 0;
  p->sendTime = 0;
  p->recvTime = 0;
  p->timeout = 0;

  sem_wait( &bufferEmpty );

  pthread_mutex_lock ( &bufferAdd );
  inFlightPackets[nextSlot % MAXWINDOW] = p;
  nextSlot ++;
  pthread_mutex_unlock ( &bufferAdd );
  
  sem_post ( &bufferFull );

  
#ifdef DEBUG
  int val;
  sem_getvalue( &bufferEmpty, &val );
  fprintf(stderr, "\nAdding Packet to Queue: Num:%d  Seq:%d  Free:%d *ACK:%d*", 
	  nextSlot-1, p->sequenceNumber, MAXWINDOW - val, p->ack );
#endif


  return;

}



int my_recv( int sock, void * bufferVoid, size_t length ) {

  char * buffer = (char *) bufferVoid;

  if (receivedFin) {
    return 0;
  }

  int numRecv = length / MAX_SEGMENT;
  int received[numRecv];
  int i;
  for (i = 0; i < numRecv; i++) {
    received[i] = 0;
  }
  int receivedSoFar = 0;
  int finReceived = 0;
  int bytesReceived = 0;

#ifdef DEBUG
  fprintf(stderr, "\nRecv:  Space-%zu, Num-%d, Base-%d", length, numRecv,
	  recvBase);
#endif



  while ( receivedSoFar < numRecv && !finReceived ) {
    
#ifdef DEBUG
    fprintf(stderr, "\nRecv: Space Remaining - %d/%d",
	    receivedSoFar, numRecv);
#endif


    struct sockaddr_in fromaddr;
    socklen_t addrlen = sizeof(fromaddr);

    // Read in a header
    char packet[ MAX_PACKET ];
    struct lab6_header * header = (struct lab6_header*) packet;
    int recv_count = recvfrom( sock, packet, MAX_PACKET,
			       0, (struct sockaddr *) &fromaddr,
			       &addrlen);
    
    if ( recv_count < 0 ) {
      // TODO - set a sender closed flag
    }
        
    if ( connect(sock, (struct sockaddr*) &fromaddr, addrlen)) {
      perror("connect (my_recv)");
    }
    
    // Sender will only send a fin once we have received all of the
    // data. So, if there is a fin, then return what we have and
    // next time we are called return 0;
    if ( ntohs(header->fin_flag) ) {
      receivedFin = 1;
      break;
    }

    // Figure out if the received packet is in our window
    int seqNum = ntohl( header->sequence_num );
    if ( seqNum < recvBase ) {
      // We've already returned the data to the user
      // so just ack it and continue.
      send_ack( sock, header );
      continue;
    }
    int slot = ( seqNum - recvBase ) / MAX_SEGMENT;

#ifdef DEBUG
    fprintf(stderr, 
	    "\nReceived Packet: Num:%d  Seq:%d  Slot:%d NumRecv:%d InWin:%d  Old:%d",
	    ntohs(header->packet_num), seqNum, slot, numRecv,
	    ( slot < numRecv ? 1 : 0 ), ( received[slot] ));
#endif

    if ( slot < numRecv ) {
      send_ack( sock, header );
    }
    else {
      // Data is out of current window
      continue;
    }
    if ( seqNum < recvBase ) {
      continue;
    }

    // Copy the data into place as needed
    if ( !received[slot] ) {
#ifdef DEBUG
      fprintf(stderr, "\nReceiver: Looks like new data. Putting it index %d",
	      seqNum - recvBase);
      fprintf(stderr, "\nData length: %d", ntohs(header->length));
      char few[10];
      strncpy(few, &packet[sizeof(struct lab6_header)], 9);
      few[9] = '\0';
      fprintf(stderr, "\nFirst bit of data is: %s", few );
#endif      
      received[slot] = 1;
      memcpy( &buffer[ seqNum - recvBase ], &packet[sizeof(struct lab6_header)],
	      ntohs(header->length) );
      receivedSoFar ++;
      bytesReceived += ntohs(header->length);
    }



  }

  recvBase += bytesReceived;

  return bytesReceived;

}

void send_ack( int sock, struct lab6_header * header ) {

  struct lab6_header ack;
  memcpy(&ack, header, sizeof(struct lab6_header));
  ack.ack_flag = htonl( 1 );
  send( sock, &ack, sizeof(ack), 0 );

#ifdef DEBUG
  fprintf(stderr, "\nSent Ack: Num:%d  Seq:%d", ntohs(ack.packet_num),
	  ntohl(ack.sequence_num));
#endif


}


void * senderFunc( void * arguments ) {

#ifdef DEBUG
  fprintf(stderr, "\nHello from sender thread!");
#endif


  uint64_t sock = (uint64_t) arguments;
  
  while ( 1 ) {

    // If there is nothing in the buffer and we are in a closing state, then
    // we want to send a fin packet.
    

    // If it has, then figure out how many things are pending
    int numPending;
    sem_getvalue ( &bufferFull, &numPending );

    if ( numPending == 0 ) {

      int localState;
      pthread_mutex_lock( &closeLock );
      localState = state;
      pthread_mutex_unlock ( &closeLock );
#ifdef DEBUG
      fprintf(stderr, "\nNothing Pending: Current State is %d", localState);
#endif
      if ( localState == NORMAL ) {
	sem_wait( &bufferFull );
	sem_post ( &bufferFull );
	numPending ++;
	// Close wasn't called yet, so don't do anything
      }
      else if ( localState == CLOSING ) {
	// Enqueue a fin packet to send
	my_send( sock, NULL, 0 );
	sem_wait( &bufferFull );
	sem_post ( &bufferFull );
	numPending ++;
      }
      else if ( localState == FIN_SENT ) {
	// Nothing pending and sender has requested a close
	pthread_exit(NULL);
      }
    }
    
    
    
    int sendNum = ( numPending > windowSize ? windowSize : numPending );

#ifdef DEBUG
    fprintf(stderr, "\nSender Thread: Num Pending-%d  Window-%d  Sending-%d",
	    numPending, windowSize, sendNum);
    /*
    int k;
    
    for ( k = 0; k < numPending; k ++ ) {
      fprintf(stderr, "\n ** Sanity Check: Packet Num %d   Seq: %d  Ack Status %d",
	      startSlot + k, 
	      inFlightPackets[ (startSlot + k) % MAXWINDOW ]->sequenceNumber,
	      inFlightPackets[ (startSlot + k) % MAXWINDOW ]->ack);
      
    }
    */
#endif

    if ( numPending != 0 ) {
      // There is data in the buffer to send
      int i;
      for ( i = 0; i < sendNum; i++ ) {
	sendPacket( sock, inFlightPackets[ (startSlot + i) % MAXWINDOW ] );
      }
      int numAcks = receiveAcks( sock, sendNum );
#ifdef DEBUG
      fprintf(stderr, "\nSender Thread: Received %d Acks", numAcks);
#endif      
      updateVariables( startSlot, sendNum );
    }
    else {
      
    }
  } /* While 1 loop */
} /* Function senderFunc */


void sendPacket( int sock, packet * p ) {

  if ( p->ack ) {
#ifdef DEBUG
    fprintf(stderr, "\nPacket Number %d already sent and ack'd!", p->sequenceNumber);
#endif
    return;
  }

  p->sendTime = current_msec();
  send( sock, p->data, p->length + sizeof(struct lab6_header), 0 );
#ifdef DEBUG
  struct lab6_header * header = (struct lab6_header*) p->data;
  fprintf(stderr, "\nSent Packet: Num-%d; Seq-%d; Ack-%d; Fin-%d", 
	  ntohs(header->packet_num),
	  ntohl(header->sequence_num), ntohl(header->ack_flag), 
	  ntohs(header->fin_flag));
#endif

}

int receiveAcks( int sock, int num ) {

  // Call select() up to num times, each time it returns, update the
  // timeout and RTT estimates, along with marking the packet as 
  // acked.
  int numReceived = 0;
  int ret;
  fd_set rfds;
  struct timeval t;
  msec_to_timeval( currentTimeout, &t );

  while ( numReceived < num ) {
    
#ifdef DEBUG_MEGA

    int numPending;
    sem_getvalue( &bufferSem, &numPending );
    numPending = MAXWINDOW - numPending;
    
    int k;
    
    for ( k = 0; k < numPending; k ++ ) {
      fprintf(stderr, "\n ** Ack Sanity Check: Packet Num %d   Seq: %d  Ack Status %d",
	      startSlot + k, 
	      inFlightPackets[ (startSlot + k) % MAXWINDOW ]->sequenceNumber,
	      inFlightPackets[ (startSlot + k) % MAXWINDOW ]->ack);
    }
#endif



    FD_ZERO(&rfds);
    FD_SET(sock, &rfds);
    
    ret = select( sock + 1, &rfds, NULL, NULL, &t );
    if (ret < 0) {
      perror("select");
      exit(1);
    }
    if ( ret == 0 ) {
      // Timeout
      return numReceived;
    }
    
    struct sockaddr_in fromaddr;
    socklen_t addrlen = sizeof(fromaddr);
    struct lab6_header receiveHeader;
    memset( &receiveHeader, 0, sizeof(struct lab6_header) );
    
#ifdef DEBUG
    fprintf(stderr, "\nGot something!!");
#endif

    int recv_count = recvfrom( sock, &receiveHeader, sizeof(struct lab6_header),
			       0, (struct sockaddr *) &fromaddr, 
			       &addrlen );
#ifdef DEBUG
    fprintf(stderr, "\nGot %d bytes", recv_count);
#endif

    if ( recv_count < (int) sizeof(struct lab6_header) ) {
      recipientClosed = 1;
      return -1;
    }
    


    if ( ntohl(receiveHeader.ack_flag) ) {
      packet * p = inFlightPackets[ ntohs(receiveHeader.packet_num) % MAXWINDOW ];
      if ( ntohs(receiveHeader.packet_num) < startSlot ) {
#ifdef DEBUG
	fprintf(stderr, "\nGot an ack for something out of window!!");
#endif	
	continue;
      }
      p->recvTime = current_msec();

#ifdef DEBUG
      fprintf(stderr, "\nPacket prior status: ack-%d", p->ack);
#endif
      p->ack = 1;

#ifdef DEBUG
      fprintf(stderr, "\nReceived Ack for packet num %d at index %d",
	      ntohs(receiveHeader.packet_num), 
	      ntohs(receiveHeader.packet_num) % MAXWINDOW );
#endif
      numReceived ++ ;
    }

    if ( ntohs(receiveHeader.fin_flag) && ntohl(receiveHeader.ack_flag) ) {
#ifdef DEBUG
      fprintf(stderr, "\nReceived an ack for Fin, and changing state.");
#endif
      pthread_mutex_lock( &closeLock );
      state = FIN_SENT;
      pthread_mutex_unlock ( &closeLock );
    }


  } /* numReceived < num */
  return numReceived;

}


void updateVariables( int startIdx, int sendNum ) {

  int i;
  int lossEvent = 0;
  int newBaseOffset = -1;
  for (i = 0; i < sendNum; i ++ ){
    packet * p = inFlightPackets[ (startIdx + i) % MAXWINDOW ];
#ifdef DEBUG
    fprintf(stderr, "\nChecking Packet %d: Ack'd-%d  TO-%d",
	    startIdx + i, p->ack, p->timeout);
#endif    
    if ( p->ack ) {
      if ( !p->timeout ) {
	updateTimeRTT( p->sendTime, p->recvTime );
      }
    }
    else {
      if ( newBaseOffset < 0 ) {
	newBaseOffset = i;
      }
      p->timeout = 1;
      lossEvent = 1;
    }
  }
  if ( lossEvent ) {
    if ( growthRate == SLOWSTART ) {
      growthRate = AIMD;
    }
    windowSize = windowSize / 2;
    currentTimeout = 2 * currentTimeout;
  }
  else {
    if ( growthRate == SLOWSTART ) {
      windowSize *= 2;
    }
    else {
      windowSize += 1;
    }
    currentTimeout = currentRTT + 4 * currentVar;
  }

  newBaseOffset = ( newBaseOffset < 0 ? sendNum : newBaseOffset );

  windowSize = ( windowSize > MAXWINDOW ? MAXWINDOW : windowSize );

  maxWindowSize = ( maxWindowSize > windowSize ?
		    maxWindowSize : windowSize );
  currentTimeout = ( currentTimeout > 3000 ? 3000 : currentTimeout );

#ifdef DEBUG
  fprintf(stderr, "\nUpdate Vars: New Window-%d  TO-%d  startSlot-%d  Shift-%d",
	  windowSize, currentTimeout, startSlot + newBaseOffset, newBaseOffset);
#endif

  for (i = 0; i < newBaseOffset; i++) {
    free ( inFlightPackets[ (startIdx + i) % MAXWINDOW ] );
    inFlightPackets[ (startIdx + i) % MAXWINDOW ] = NULL;
    sem_post( &bufferEmpty );
    sem_wait( &bufferFull );
  }
  startSlot += newBaseOffset;
  
}

void updateTimeRTT( int start, int end ) {

  double cRTT = (0.875) * currentRTT + (0.125) * (end - start);
  double cVar = (0.75) * currentVar + 
    ( end - start > currentRTT ? 0.25 * ( end-start - currentRTT ) :
      0.25 * ( currentRTT - ( end - start) ));
  currentRTT = (int) cRTT;
  currentVar = (int) cVar;

}


int my_close( int sock ) {


  if ( receivedFin ) {
    fd_set rfds;
    int select_result = 1;
    
#ifdef DEBUF
    fprintf(stderr, "\nEntering time wait for %d ms.\n", TIME_TO_WAIT);
#endif
    while ( select_result > 0 ) {
      FD_ZERO(&rfds);
      FD_SET(sock, &rfds);
      struct timeval timeout;
      char fin_packet[sizeof(struct lab6_header)];
      memset(fin_packet, 0, sizeof(struct lab6_header));

      msec_to_timeval(TIME_TO_WAIT, &timeout);
      select_result = select(sock+1, &rfds, NULL, NULL, &timeout);
      
      if ( 1 == select_result ) {
	recv(sock, fin_packet, sizeof(struct lab6_header), 0);
	struct lab6_header * hdr = (struct lab6_header*) fin_packet;
	if (!hdr->fin_flag) {
	  fprintf(stderr, "\nGot a non-fin packet in close?\n");
	}
	else {
	  fprintf(stderr, 
		  "\nGot anothher fin - resetting timer to %d ms.",
		  TIME_TO_WAIT);
	}
	send_ack(sock, hdr);
      }
    }
    
    pthread_mutex_destroy( &closeLock );
    pthread_mutex_destroy( &bufferAdd );
    sem_destroy( &bufferEmpty );
    sem_destroy( &bufferFull );

    return close(sock);
  }


  //my_send( sock, NULL, 0 );
  pthread_mutex_lock ( &closeLock );
  state = CLOSING;
  pthread_mutex_unlock ( &closeLock );
  pthread_join( senderThread, NULL );




  pthread_mutex_destroy( &closeLock );
  pthread_mutex_destroy( &bufferAdd );
  sem_destroy( &bufferEmpty );
  sem_destroy( &bufferFull );
  

  return close(sock);

}




/* The following three helper functions make dealing with time values a bit
 * easier.  You probably don't need to modify them, but you can if you want. */
int timeval_to_msec(struct timeval *t) {
    return t->tv_sec*1000+t->tv_usec/1000;
}

void msec_to_timeval(int millis, struct timeval *out_timeval) {
    out_timeval->tv_sec = millis/1000;
    out_timeval->tv_usec = (millis%1000)*1000;
}

int current_msec() {
    struct timeval t;
    gettimeofday(&t,0);
    return timeval_to_msec(&t);
}
/* End time functions. */


int my_connect(int socket,struct sockaddr_in *toaddr,int addrsize) {
  closed = 0;
  uint64_t sock = socket;
  //pthread_create( &senderThread, NULL, senderFunc, (void *) sock );
  return connect(socket,(struct sockaddr*)toaddr,addrsize);
  return 0;
}

int my_rtt(int sock) {
    return currentRTT;
}

/* The value this returns (max_window) should be the value of the largest
 * window size your library uses over the course of the run. */
int my_max_window(int sock) {
  return maxWindowSize;
}

void * Malloc ( size_t size ) {

  void * toRet = malloc(size);
  if ( ! toRet ) {
    perror("malloc");
    exit(1);
  }
  memset ( toRet, 0, size );
  return toRet;


}
